load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_csm.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/shea_util.ncl"
;**************************************
begin

season = getenv("season")
VAR_measure = getenv("VAR_measure")
ilevel = toint(getenv("ilevel"))

season = "JJA"
VAR_measure = "VAR_R95_R50"
ilevel = 4

VAR_measure_regional = VAR_measure+"_regional"


;---------------------

Nmodel = 60
N1 = 35        ;CMIP5
N2 = 25        ;CMIP6


nlat = 180
nlon = 360

Nindex = 4
index_name = (/"R90","R95","R98","R99"/)

Nvar = 3
var_name = (/"pr1d","pr5d","pr30d"/)

Nlevel = 6
warming_levels = (/"0C","0.5degC","1degC","2degC","3degC","4degC"/)

Nobs = 4
dataset_obs = (/"GPCP","GPCC","CPC ","ERA5"/)

NSMILEs = 5
SMILEs_model = (/"CESM1-CAM5","CSIRO-Mk3-6-0","CanESM2","EC-EARTH","GFDL-CM3"/)


;*******************************************
;regional
;*******************************************
lonW := (/0,0,80,115,0,260/)
lonE := (/360,360,150,180,25,300/)
latS := (/45,-70,45,30,50,50/)
latN := (/70,-45,70,45,70,70/)

subregions := (/"NH_extratropics","SH_extratropics","Northern_Asia","WNP_monsoon","Europe","Eastern_NAmerica"/)
subregions_plot := (/"NH mid-latitudes","SH mid-latitudes","Northern Asia","West North Pacific","Europe","Eastern North America"/)
domain := (/"ocean","ocean","land","ocean","land","land"/)
Nregion := dimsizes(lonW)



;--------------read change in frq---------------
diri := "/WORK/zhangwx/EConstraint/data/"
change := new((/Nvar,Nmodel,Nindex,nlat,nlon/),float)

do ivar = 0,Nvar-1

  f1 := addfile(diri+"CMIP5/pr/"+var_name(ivar)+"/"+season+"/RX_pct/multimodel/change_frequency_RX_pct_"+var_name(ivar)+"_multimodel_warming_levels_"+season+".nc","r")
  change(ivar,0:N1-1,:,:,:) = f1->change_frq_RX_pct(:,ilevel,:,:,:)               ;(35model, 4index, lat, lon)

  f2 := addfile(diri+"CMIP6/pr/"+var_name(ivar)+"/"+season+"/RX_pct/multimodel/change_frequency_RX_pct_"+var_name(ivar)+"_multimodel_warming_levels_"+season+".nc","r")
  change(ivar,N1:N1+N2-1,:,:,:) = f2->change_frq_RX_pct(:,ilevel,:,:,:)        ;(25model 4index lat lon)

end do

PR := change/100.+1             ;% change --> probability ratio
copy_VarCoords(change,PR)


;-------------read model pd Variability-------------
variability := new((/Nvar,Nmodel,nlat,nlon/),float)

do ivar = 0,Nvar-1

  f1 := addfile(diri+"CMIP5/pr/"+var_name(ivar)+"/"+season+"/variability/multimodel/variability_"+var_name(ivar)+"_multimodel_"+season+"_1997-2014.nc","r")
  variability(ivar,0:N1-1,:,:) = f1->$VAR_measure$          ;(35model lat lon)

  f2 := addfile(diri+"CMIP6/pr/"+var_name(ivar)+"/"+season+"/variability/multimodel/variability_"+var_name(ivar)+"_multimodel_"+season+"_1997-2014.nc","r")
  variability(ivar,N1:N1+N2-1,:,:) = f2->$VAR_measure$       ;(25model lat lon)

end do


;----------------read model warming periods-----------------
warming_periods := new(Nmodel,integer)

f1 := addfile(diri+"CMIP5/tas/GMST/multimodel/GMST_multimodel_warming_levels_year.nc","r")
warming_periods(0:N1-1) = f1->year_center(:,ilevel)                   ;(35model)

f2 := addfile(diri+"CMIP6/tas/GMST/multimodel/GMST_multimodel_warming_levels_year.nc","r")
warming_periods(N1:N1+N2-1) = f2->year_center(:,ilevel)               ;(25model)


Nmodel_valid := num(.not.(ismissing(warming_periods)))


;**************read OBS variability******************
VAR_obs := new((/Nvar,Nobs,nlat,nlon/),float)

do ivar = 0,Nvar-1
  fobs := addfile(diri+"OBS/MultiOBS/"+var_name(ivar)+"/variability/variability_"+var_name(ivar)+"_MultiOBS_"+season+"_1997-2014.nc","r")
  VAR_obs(ivar,:,:,:) = fobs->$VAR_measure$               ;(4dataset lat lon)
end do



;***************************************
;mask to land
;***************************************
a := addfile("$NCARG_ROOT/lib/ncarg/data/cdf/landsea.nc","r")
lsdata := a->LSMASK

lsm := landsea_mask(lsdata,PR&latitude,PR&longitude)
lsm1 := conform(PR,lsm,(/3,4/))                 ;(/Nvar,Nmodel,Nindex,lat,lon/)

PR_land := where(((lsm1 .eq. 1) .or. (lsm1 .eq. 3)), PR, PR@_FillValue)
copy_VarCoords(PR,PR_land)

variability_land := where(((lsm1(:,:,0,:,:) .eq. 1) .or. (lsm1(:,:,0,:,:) .eq. 3)), variability, variability@_FillValue)
copy_VarCoords(variability,variability_land)

VAR_obs_land := where(((lsm1(:,0:Nobs-1,0,:,:) .eq. 1) .or. (lsm1(:,0:Nobs-1,0,:,:) .eq. 3)), VAR_obs, VAR_obs@_FillValue)
copy_VarCoords(VAR_obs,VAR_obs_land)


;***********************************************
;regional mean
;***********************************************
lat := variability&latitude
rad = 4*atan(1.0)/180.
clat := cos(lat*rad)
clat!0 = "lat"
clat&lat = lat

variability_regional := new((/Nvar,Nmodel,Nregion/),float)
PR_regional := new((/Nvar,Nmodel,Nindex,Nregion/),float)
VAR_obs_regional := new((/Nvar,Nobs,Nregion/),float)

do iregion = 0,Nregion-1
  variability_regional(:,:,iregion) = wgt_areaave_Wrap(variability(:,:,{latS(iregion):latN(iregion)},{lonW(iregion):lonE(iregion)}),clat({latS(iregion):latN(iregion)}),1.0,0)
  PR_regional(:,:,:,iregion) = wgt_areaave_Wrap(PR(:,:,:,{latS(iregion):latN(iregion)},{lonW(iregion):lonE(iregion)}),clat({latS(iregion):latN(iregion)}),1.0,0)
  VAR_obs_regional(:,:,iregion) = wgt_areaave_Wrap(VAR_obs(:,:,{latS(iregion):latN(iregion)},{lonW(iregion):lonE(iregion)}),clat({latS(iregion):latN(iregion)}),1.0,0)

  if (domain(iregion).eq. "ocean") then
    VAR_obs_regional(:,1:2,iregion) = VAR_obs_regional@_FillValue
  end if            ;set GPCC & CPC to missing for ocean domains

end do


;-------------for land regions 4-5-----------
do iregion = 4,5                       ;Northern Europe & Northeast America
  variability_regional(:,:,iregion) = wgt_areaave_Wrap(variability_land(:,:,{latS(iregion):latN(iregion)},{lonW(iregion):lonE(iregion)}),clat({latS(iregion):latN(iregion)}),1.0,0)
  PR_regional(:,:,:,iregion) = wgt_areaave_Wrap(PR_land(:,:,:,{latS(iregion):latN(iregion)},{lonW(iregion):lonE(iregion)}),clat({latS(iregion):latN(iregion)}),1.0,0)
  VAR_obs_regional(:,:,iregion) = wgt_areaave_Wrap(VAR_obs_land(:,:,{latS(iregion):latN(iregion)},{lonW(iregion):lonE(iregion)}),clat({latS(iregion):latN(iregion)}),1.0,0)
end do


;-------------------------------
VAR_obs_regional_MME := dim_avg_n_Wrap(VAR_obs_regional,1)            ;(/Nvar,Nregion/)

;print(VAR_obs_regional(1,:,4))


;===========linear fit===========
fit = new((/Nvar,Nmodel,Nindex,Nregion/),float)
corr = new((/Nvar,Nindex,Nregion/),float)    ;Spearman correlation coefficient
slope = new((/Nvar,Nindex,Nregion/),float)

do ivar = 0,Nvar-1
  do iindex = 0,Nindex-1
    do iregion = 0,Nregion-1

      rc := regline_stats( variability_regional(ivar,:,iregion), PR_regional(ivar,:,iindex,iregion) )
      b := rc@b
      fit(ivar,:,iindex,iregion) = b(1)*variability_regional(ivar,:,iregion)+b(0)
      slope(ivar,iindex,iregion) = b(1)

      corr(ivar,iindex,iregion) = spcorr( variability_regional(ivar,:,iregion), PR_regional(ivar,:,iindex,iregion) )

    end do
  end do
end do

prob = rtest(corr,Nmodel_valid,0)       ;significance level ;<0.10/0.05 means significant



;*************************************************************************************************
;read internal variability from SMILEs
;*************************************************************************************************
IV_SMILEs_best := new((/NSMILEs,Nvar,Nregion/),float)
IV_SMILEs_stddev := new((/NSMILEs,Nvar,Nregion/),float)
IV_SMILEs_min := IV_SMILEs_best
IV_SMILEs_max := IV_SMILEs_best

do is = 0,NSMILEs-1
  do ivar = 0,Nvar-1

    diriv := diri+"SMILEs/"+SMILEs_model(is)+"/pr/pd/"+var_name(ivar)+"/"+season+"/variability/multirun/"
    fiv := addfile(diriv+"variability_"+var_name(ivar)+"_"+SMILEs_model(is)+"_multirun_1997-2014_"+season+"_regional.nc","r")
    pdVAR_IV_tmp := fiv->$VAR_measure_regional$            ;(runs, region)


    IV_SMILEs_best(is,ivar,:) = dim_avg_n_Wrap(pdVAR_IV_tmp,0)           ;member-mean
    IV_SMILEs_stddev(is,ivar,:) = dim_stddev_n_Wrap(pdVAR_IV_tmp,0)      ;intermember stddev
    IV_SMILEs_min(is,ivar,:) = dim_min_n_Wrap(pdVAR_IV_tmp,0)
    IV_SMILEs_max(is,ivar,:) = dim_max_n_Wrap(pdVAR_IV_tmp,0)

  end do
end do


IV_SMILEs_stddev_MME := dim_avg_n_Wrap(IV_SMILEs_stddev,0)           ;5-SMILEs mean internal variability (stddev)   (/Nvar,Nregion/)



;***************************************************************
;do ivar = 0,Nvar-1
do ivar = 0,0                        ;pr1d
do iindex = 1,1                      ;R95
;do iindex = 2,3               ;R98 R99

wks := gsn_open_wks("x11","/WORK/zhangwx/EConstraint/pic/correlation/onlyCMIP/frequency/Scatter_Variability_PR_regions_"+warming_levels(ilevel)+"_"+VAR_measure+"_"+var_name(ivar)+"_"+index_name(iindex)+"_"+season)

plot := new((/Nregion/),graphic)
plot_pl := plot
plot_fit := plot
plot_corr := plot
plot_pm := new((/Nregion,Nmodel/),graphic)
plot_obs := new((/Nregion,Nobs/),graphic)
plot_obs_MME := new((/Nregion/),graphic)

plot_tx := new((/Nregion,2/),graphic)

plot_IV_SMILEs_best := new((/Nregion,NSMILEs/),graphic)
plot_IV_SMILEs_stddev := new((/Nregion,NSMILEs/),graphic)
plot_IV_stddev_MME := new((/Nregion/),graphic)


res = True
res@gsnFrame = False
res@gsnDraw = False

res@tmXTOn = False
res@tmYROn = False

;res@tmXBMajorOutwardLengthF = 0.
;res@tmXBMinorOutwardLengthF = 0.
;res@tmYLMajorOutwardLengthF = 0.
;res@tmYLMinorOutwardLengthF = 0.

res@tmBorderThicknessF = 1.0
res@tmXBMajorThicknessF = 1.0
res@tmXBMinorThicknessF = 1.0
res@tmYLMajorThicknessF = 1.0
res@tmYLMinorThicknessF = 1.0


res@xyMarkLineModes = "Markers"
res@xyMarkers =  16
res@xyMarkerColor = "white"
res@xyMarkerSizeF = 0.01

;res@tiMainFontHeightF = 0.034
res@tiXAxisFontHeightF = 0.023
res@tiYAxisFontHeightF = 0.023
res@tmXBLabelFontHeightF = 0.02
res@tmYLLabelFontHeightF = 0.02
res@gsnLeftStringFontHeightF = 0.026
res@gsnRightStringFontHeightF = 0.026
;res@gsnRightStringOrthogonalPosF = -0.15
;res@gsnRightStringParallelPosF = 0.93


;res@xyLineThicknessF = 1.
;res@xyLineColor = "grey"


;-------------------------

plres = True
plres@gsLineThicknessF = 2.
plres@gsLineColor = "grey70"
plres@gsLineDashPattern = 0


;-------------------------
colors := (/"cornflowerblue","brown1","orange"/)
markercolor = new(Nmodel,string)
markercolor(0:N1-1) = colors(0)
markercolor(N1:N1+N2-1) = colors(1)

modelcolor = new(Nmodel,string)
modelcolor(0:N1-1) = colors(0)
modelcolor(N1:N1+N2-1) = colors(1)

pmres = True
pmres@gsMarkerIndex =  16
pmres@gsMarkerSizeF = 0.005

;------------------------
txres = True
txres@txFontHeightF = 0.022
txres@txFontColor = "black"

txres_model = True
txres_model@txFontHeightF = 0.017


;*****************************************
siglvl = 0.05

leftst = (/"a","b","c","d","e","f"/)

YMin = new((/Nregion/),float)
YMax = YMin
XMin = YMin
XMax = YMin
height = YMin
width = YMin


do iregion = 0,Nregion-1

  res@tiXAxisString = "Present-day variability (mm/day)"
  res@tiYAxisString = "Probability ratio"


  height(iregion) = (max(PR_regional(ivar,:,iindex,iregion)) - min(PR_regional(ivar,:,iindex,iregion)))*1.3
  YMin(iregion) = min(PR_regional(ivar,:,iindex,iregion)) - height(iregion)*0.15
  YMax(iregion) = max(PR_regional(ivar,:,iindex,iregion)) + height(iregion)*0.15

  width(iregion) = (max(variability_regional(ivar,:,iregion)) - min(variability_regional(ivar,:,iregion)))*1.2
  XMin(iregion) = min(variability_regional(ivar,:,iregion)) - width(iregion)*0.1
  XMax(iregion) = max( (/max(variability_regional(ivar,:,iregion)),max(VAR_obs_regional(ivar,:,iregion)) /)) + width(iregion)*0.1

  res@trXMinF = XMin(iregion)
  res@trXMaxF = XMax(iregion)
  res@trYMinF = YMin(iregion)
  res@trYMaxF = YMax(iregion)

  res@gsnLeftString = subregions_plot(iregion)
  res@gsnRightString = season
  plot(iregion) = gsn_csm_xy(wks,variability_regional(ivar,:,iregion),PR_regional(ivar,:,iindex,iregion),res)

  if (prob(ivar,iindex,iregion).le. siglvl) then            ;significant correlation
    plot_fit(iregion) = gsn_add_polyline(wks,plot(iregion),variability_regional(ivar,:,iregion),fit(ivar,:,iindex,iregion),plres)
  end if

  tx_corr := "r="+sprintf("%4.2f",corr(ivar,iindex,iregion))
  if (prob(ivar,iindex,iregion).le. siglvl) then            ;significant correlation
    tx_corr = tx_corr+"*"
  end if
  tx_slope := "slope="+sprintf("%4.2f",slope(ivar,iindex,iregion))

  plot_tx(iregion,0) = gsn_add_text(wks,plot(iregion),tx_corr,XMax(iregion)-0.14*width(iregion),YMax(iregion)-0.12*height(iregion),txres)
  plot_tx(iregion,1) = gsn_add_text(wks,plot(iregion),tx_slope,XMax(iregion)-0.19*width(iregion),YMax(iregion)-0.2*height(iregion),txres)


  do imodel = 0,Nmodel-1
    if .not.(ismissing(warming_periods(imodel))) then

      txres_model@txFontColor = modelcolor(imodel)
      plot_pm(iregion,imodel) = gsn_add_text(wks,plot(iregion),tostring(imodel+1),variability_regional(ivar,imodel,iregion),PR_regional(ivar,imodel,iindex,iregion),txres_model)

      ;pmres@gsMarkerColor = markercolor(imodel)
      ;plot_pm(iregion,imodel) = gsn_add_polymarker(wks,plot(iregion),variability_regional(ivar,imodel,iregion),PR_regional(ivar,imodel,iindex,iregion),pmres)
    end if
  end do

end do


;----------------obs------------------
plres@gsLineThicknessF = 2.5        ;obs MME
plres@gsLineColor = "grey70"
plres@gsLineDashPattern = 0

pmres@gsMarkerSizeF = 0.011         ;individual obs
pmres@gsMarkerColor = "grey40"

marker_obs = (/12,7,2,9/)

do iregion = 0,Nregion-1

  do iobs = 0,Nobs-1
    if (.not.(ismissing(VAR_obs_regional(ivar,iobs,iregion)))) then
    pmres@gsMarkerIndex = marker_obs(iobs)
    plot_obs(iregion,iobs) = gsn_add_polymarker(wks,plot(iregion),VAR_obs_regional(ivar,iobs,iregion),(YMax(iregion)-height(iregion)*0.05),pmres)
    end if
  end do

  plot_obs_MME(iregion) = gsn_add_polyline(wks,plot(iregion),(/VAR_obs_regional_MME(ivar,iregion),VAR_obs_regional_MME(ivar,iregion)/),(/-5000,5000/),plres)

end do




;--------------------------
txres = True
txres@txAngleF = 0
txres@txFontHeightF = 0.017

;gsn_text_ndc(wks,var_name,fspan(0.26,0.78,3),fspan(0.97,0.97,3),txres)

txres@txAngleF = 90
;gsn_text_ndc(wks,index_name,fspan(0.08,0.08,4),fspan(0.86,0.14,4),txres)


;-------------
xx = fspan(0.29,0.29,Nobs)
yy = fspan(0.75,0.7,Nobs)
txres@txFontHeightF = 0.008
txres@txAngleF = 0
txres@txFontColor = "black"

do iobs = 0,Nobs-1
  pmres@gsMarkerIndex = marker_obs(iobs)
  gsn_polymarker_ndc(wks,xx(iobs)-0.03,yy(iobs),pmres)
  gsn_text_ndc(wks,dataset_obs(iobs),xx(iobs),yy(iobs),txres)
end do


;**********************internal variability************************
pmres_IV = True
pmres_IV@gsMarkerIndex = 16
pmres_IV@gsMarkerSizeF = 0.002
pmres_IV@gsMarkerColor = "grey"

plres_IV = True
plres_IV@gsLineColor = "grey60"
plres_IV@gsLineDashPattern = 0
plres_IV@gsLineThicknessF = 2.

yy_IV = fspan(-5,5,NSMILEs)

do iregion = 0,Nregion-1

    yy_IV := fspan(YMin(iregion)+0.02*height(iregion),YMin(iregion)+0.12*height(iregion), NSMILEs)

    do is = 0,NSMILEs-1
      plot_IV_SMILEs_best(iregion,is) = gsn_add_polymarker(wks,plot(iregion),IV_SMILEs_best(is,ivar,iregion),yy_IV(is),pmres_IV)                       ;(/NSMILEs,Nvar,Nregion/)

      x1 := IV_SMILEs_best(is,ivar,iregion)-IV_SMILEs_stddev(is,ivar,iregion)
      x2 := IV_SMILEs_best(is,ivar,iregion)+IV_SMILEs_stddev(is,ivar,iregion)
;      x1 := IV_SMILEs_min(is,ivar,iregion)
;      x2 := IV_SMILEs_max(is,ivar,iregion)
      plot_IV_SMILEs_stddev(iregion,is) = gsn_add_polyline(wks,plot(iregion),(/x1,x2/),(/yy_IV(is),yy_IV(is)/),plres_IV)
    end do

end do


;------------
txres2 = True
txres2@txFont = "helvetica-bold"
txres2@txFontHeightF = 0.012
gsn_text_ndc(wks,leftst(0:5:2),fspan(0.2,0.2,3),fspan(0.94,0.31,3),txres2)
gsn_text_ndc(wks,leftst(1:5:2),fspan(0.52,0.52,3),fspan(0.94,0.31,3),txres2)


;--------------------
resP = True
resP@gsnPanelTop = 0.95
;resP@gsnPanelLeft = 0.05
;resP@gsnPanelLabelBar  = True
resP@gsnPanelYWhiteSpacePercent = 3.
resP@gsnPanelXWhiteSpacePercent = 4.
;resP@lbLabelFontHeightF = 0.012

gsn_panel(wks,plot,(/3,2/),resP)


end do                  ;{iindex}
end do                  ;{ivar}


end
